""""""
from datetime import datetime
from typing import List

from pymongo import ASCENDING, MongoClient, ReplaceOne
from pymongo.database import Database
from pymongo.cursor import Cursor
from pymongo.collection import Collection
from pymongo.results import DeleteResult

from vnpy.trader.constant import Exchange, Interval
from vnpy.trader.object import BarData, TickData, FutBasicData, FutMappingData
from vnpy.feeds.database import BaseDatabase, BarOverview, DB_TZ

from pandas import DataFrame
import pandas as pd


class MongodbDatabase(BaseDatabase):
    """MongoDB数据库接口"""

    def __init__(self, database: str = "tushare", host: str = "localhost", port: int = 27017,
                 username: str = "root", password: str = "") -> None:
        """"""
        self.database: str = database
        self.host: str = host
        self.port: int = port
        self.username: str = username
        self.password: str = password

        # 创建客户端
        if self.username and self.password:
            self.client: MongoClient = MongoClient(
                host=self.host,
                port=self.port,
                tz_aware=True,
                username=self.username,
                password=self.password,
                tzinfo=DB_TZ
            )
        else:
            self.client: MongoClient = MongoClient(
                host=self.host,
                port=self.port,
                tz_aware=True,
                tzinfo=DB_TZ
            )

        # 初始化数据库
        self.db: Database = self.client[self.database]

        # 初始化K线数据表
        self.bar_collection: Collection = self.db["bar_data"]
        self.bar_collection.create_index(
            [
                ("Exchange", ASCENDING),
                ("Symbol", ASCENDING),
                ("Interval", ASCENDING),
                ("Datetime", ASCENDING),
            ],
            unique=True
        )

        # 初始化Tick数据表
        self.tick_collection: Collection = self.db["tick_data"]
        self.tick_collection.create_index(
            [
                ("exchange", ASCENDING),
                ("symbol", ASCENDING),
                ("datetime", ASCENDING),
            ],
            unique=True
        )

        # 初始化K线概览表
        self.overview_collection: Collection = self.db["bar_overview"]
        self.overview_collection.create_index(
            [
                ("exchange", ASCENDING),
                ("symbol", ASCENDING),
                ("interval", ASCENDING),
            ],
            unique=True
        )

        # 初始化fut_basic数据表
        self.fut_basic_collection: Collection = self.db["fut_basic"]
        self.fut_basic_collection.create_index(
            [
                ("Exchange", ASCENDING),
                ("Symbol", ASCENDING),
            ],
            unique=True
        )

        # 初始化fut_mapping数据表
        self.fut_mapping_collection: Collection = self.db["fut_mapping"]
        self.fut_mapping_collection.create_index(
            [
                ("ts_code", ASCENDING),
                ("trade_date", ASCENDING),
            ],
            unique=True
        )

        # 初始化securities数据表
        self.securities_collection: Collection = self.db["securities"]
        self.securities_collection.create_index(
            [
                ("name", ASCENDING),
                ("Symbol", ASCENDING),
            ],
            unique=True
        )

    def save_bar_data(self, bars: DataFrame) -> bool:
        """保存K线数据"""
        requests: List[ReplaceOne] = []

        for bar in bars.itertuples():
            # 逐个插入
            filter: dict = {
                "Symbol": bar.Symbol,
                "Exchange": bar.Exchange.value,
                "Datetime": bar.Datetime,
                "Interval": bar.Interval.value,
            }

            d: dict = {
                "Symbol": bar.Symbol,
                "Exchange": bar.Exchange.value,
                "Datetime": bar.Datetime,
                "Interval": bar.Interval.value,
                "Volume": bar.Volume,
                "OpenInterest": bar.OpenInterest,
                "Open": bar.Open,
                "High": bar.High,
                "Low": bar.Low,
                "Close": bar.Close,
            }
            requests.append(ReplaceOne(filter, d, upsert=True))

        # bar_data = json.loads(bars.T.to_json()).values()
        # result = self.bar_collection.insert_many(bar_data)

        self.bar_collection.bulk_write(requests, ordered=False)

        # 更新汇总
        filter: dict = {
            "symbol": bar.Symbol,
            "exchange": bar.Exchange.value,
            "interval": bar.Interval.value
        }

        overview: dict = self.overview_collection.find_one(filter)
        print(bars['Datetime'].min())
        # datetime.timezone()
        # print(overview["start"])

        if not overview:
            overview = {
                "symbol": bar.Symbol,
                "exchange": bar.Exchange.value,
                "interval": bar.Interval.value,
                "count": len(bars),
                "start": bars['Datetime'].min(),
                "end": bars['Datetime'].max()
            }
        else:
            # timezone('Asia/Shanghai')
            overview["start"] = min(bars['Datetime'].min(), overview["start"].replace(tzinfo=None))
            overview["end"] = max(bars['Datetime'].max(), overview["end"].replace(tzinfo=None))
            overview["count"] = self.bar_collection.count_documents(filter)

        self.overview_collection.update_one(filter, {"$set": overview}, upsert=True)

        return True

    def save_tick_data(self, ticks: List[TickData]) -> bool:
        """保存TICK数据"""
        requests: List[ReplaceOne] = []

        for tick in ticks:
            filter: dict = {
                "symbol": tick.symbol,
                "exchange": tick.exchange.value,
                "datetime": tick.datetime,
            }

            d: dict = {
                "symbol": tick.symbol,
                "exchange": tick.exchange.value,
                "datetime": tick.datetime,
                "name": tick.name,
                "volume": tick.volume,
                "turnover": tick.turnover,
                "open_interest": tick.open_interest,
                "last_price": tick.last_price,
                "last_volume": tick.last_volume,
                "limit_up": tick.limit_up,
                "limit_down": tick.limit_down,
                "open_price": tick.open_price,
                "high_price": tick.high_price,
                "low_price": tick.low_price,
                "pre_close": tick.pre_close,
                "bid_price_1": tick.bid_price_1,
                "bid_price_2": tick.bid_price_2,
                "bid_price_3": tick.bid_price_3,
                "bid_price_4": tick.bid_price_4,
                "bid_price_5": tick.bid_price_5,
                "ask_price_1": tick.ask_price_1,
                "ask_price_2": tick.ask_price_2,
                "ask_price_3": tick.ask_price_3,
                "ask_price_4": tick.ask_price_4,
                "ask_price_5": tick.ask_price_5,
                "bid_volume_1": tick.bid_volume_1,
                "bid_volume_2": tick.bid_volume_2,
                "bid_volume_3": tick.bid_volume_3,
                "bid_volume_4": tick.bid_volume_4,
                "bid_volume_5": tick.bid_volume_5,
                "ask_volume_1": tick.ask_volume_1,
                "ask_volume_2": tick.ask_volume_2,
                "ask_volume_3": tick.ask_volume_3,
                "ask_volume_4": tick.ask_volume_4,
                "ask_volume_5": tick.ask_volume_5,
                "localtime": tick.localtime,
            }

            requests.append(ReplaceOne(filter, d, upsert=True))

        self.tick_collection.bulk_write(requests, ordered=False)

        return True

    def save_fut_basic(self, futs: DataFrame) -> bool:
        """保存K线数据"""
        requests: List[ReplaceOne] = []

        for fut in futs:
            # 逐个插入
            filter: dict = {
                "symbol": fut.symbol,
                "exchange": fut.exchange.value,
            }

            d: dict = {
                "symbol": fut.symbol,  # 合约代码
                "exchange": fut.exchange.value,  # 交易市场
                "ts_code": fut.ts_code,  # 合约代码
                "name": fut.name,  # 中文简称
                "fut_code": fut.fut_code,  # 合约产品代码
                "multiplier": fut.multiplier,  # 合约乘数
                "trade_unit": fut.trade_unit,  # 交易计量单位
                "per_unit": fut.per_unit,  # 交易单位(每手)
                "quote_unit": fut.quote_unit,  # 报价单位
                "quote_unit_desc": fut.quote_unit_desc,  # 最小报价单位说明
                "d_mode_desc": fut.d_mode_desc,  # 交割方式说明
                "list_date": fut.list_date,  # 上市日期
                "delist_date": fut.delist_date,  # 最后交易日期
                "d_month": fut.d_month,  # 交割月份
                "last_ddate": fut.last_ddate,  # 最后交割日
                "trade_time_desc": fut.trade_time_desc,  # 交易时间说明
            }

            requests.append(ReplaceOne(filter, d, upsert=True))

        self.fut_basic_collection.bulk_write(requests, ordered=False)

        return True

    def save_fut_mapping(self, futs: List[FutMappingData]) -> bool:
        """保存K线数据"""
        requests: List[ReplaceOne] = []

        for fut in futs:
            # 逐个插入
            filter: dict = {
                "ts_code": fut.ts_code,
                "trade_date": fut.trade_date,
            }

            d: dict = {
                "ts_code": fut.ts_code,  # 合约代码
                "trade_date": fut.trade_date,  # 交易日期
                "mapping_ts_code": fut.mapping_ts_code,  # 期货合约代码
            }

            requests.append(ReplaceOne(filter, d, upsert=True))

        self.fut_mapping_collection.bulk_write(requests, ordered=False)

        return True

    def load_bar_data(
            self,
            symbol: str,
            exchange: Exchange,
            interval: Interval,
            start: datetime,
            end: datetime
    ) -> DataFrame:
        """读取K线数据"""
        filter: dict = {
            "Symbol": symbol,
            "Exchange": exchange.value,
            "Interval": interval.value,
            "Datetime": {
                "$gte": start.astimezone(None),
                "$lte": end.astimezone(None)
            }
        }

        bars = DataFrame(list(self.bar_collection.find(filter)))
        bars['Datetime'] = bars['Datetime'].dt.tz_convert(None)
        bars.set_index('Datetime', inplace=True)
        bars = bars[['Open', 'High', 'Low', 'Close', 'Volume', 'OpenInterest']]
        return bars

    def load_tick_data(
            self,
            symbol: str,
            exchange: Exchange,
            start: datetime,
            end: datetime
    ) -> List[TickData]:
        """读取TICK数据"""
        filter: dict = {
            "symbol": symbol,
            "exchange": exchange.value,
            "datetime": {
                "$gte": start.astimezone(DB_TZ),
                "$lte": end.astimezone(DB_TZ)
            }
        }

        c: Cursor = self.tick_collection.find(filter)

        ticks: List[TickData] = []
        for d in c:
            d["exchange"] = Exchange(d["exchange"])
            d["gateway_name"] = "DB"
            d.pop("_id")

            tick: TickData = TickData(**d)
            ticks.append(tick)

        return ticks

    def delete_bar_data(
            self,
            symbol: str,
            exchange: Exchange,
            interval: Interval
    ) -> int:
        """删除K线数据"""
        filter: dict = {
            "symbol": symbol,
            "exchange": exchange.value,
            "interval": interval.value,
        }

        result: DeleteResult = self.bar_collection.delete_many(filter)
        self.overview_collection.delete_one(filter)

        return result.deleted_count

    def delete_tick_data(
            self,
            symbol: str,
            exchange: Exchange
    ) -> int:
        """删除TICK数据"""
        filter: dict = {
            "symbol": symbol,
            "exchange": exchange.value
        }

        result: DeleteResult = self.tick_collection.delete_many(filter)
        return result.deleted_count

    def get_bar_overview(self) -> List[BarOverview]:
        """查询数据库中的K线汇总信息"""
        c: Cursor = self.overview_collection.find()

        overviews: List[BarOverview] = []
        for d in c:
            d["exchange"] = Exchange(d["exchange"])
            d["interval"] = Interval(d["interval"])
            d.pop("_id")

            overview: BarOverview = BarOverview(**d)
            overviews.append(overview)

        return overviews

    def quary_fut_basic(self, fut_code: str = "SA", fut_type: str = "2") -> List[FutBasicData]:
        """
        Query fut basic from datafeed.
        """
        # fut_type 0，所有 1:主力 2：连续 3：普通
        filter: dict = {
            "fut_code": fut_code,
            # "name": {'$regex': '主力'}
        }
        if (fut_type == "1"):
            filter: dict = {
                "fut_code": fut_code,
                "name": {'$regex': '主力'}
            }
        elif (fut_type == "2"):
            filter: dict = {
                "fut_code": fut_code,
                "name": {'$regex': '连续'}
            }
        else:
            filter: dict = {
                "fut_code": fut_code,
            }

        c: Cursor = self.fut_basic_collection.find(filter)

        datas: List[FutBasicData] = []
        for d in c:
            d["exchange"] = Exchange(d["exchange"])
            # d["interval"] = Interval(d["interval"])
            d["gateway_name"] = "DB"
            d.pop("_id")

            data = FutBasicData(**d)
            datas.append(data)

        return datas

    def save_all_securities(self, securities) -> bool:
        """保存K线数据"""
        requests: List[ReplaceOne] = []

        for security in securities.itertuples():
            # 逐个插入
            filter: dict = {
                "Symbol": security.Symbol,
                # "Exchange": security.Exchange.value,
                "name": security.name,
            }
            # , display_name, name, start_date, end_date, type
            d: dict = {
                "Symbol": security.Symbol,
                "display_name": security.display_name,
                "name": security.name,
                "start_date": security.start_date.value,
                "end_date": security.end_date,
                "type": security.type
            }
            requests.append(ReplaceOne(filter, d, upsert=True))

        self.securities_collection.bulk_write(requests, ordered=False)
        return True

    def get_all_securities(
            self
    ) -> DataFrame:
        """读取K线数据"""
        filter: dict = {
            "Symbol": "",
                # "Exchange": security.Exchange.value,
                "name": "",
            }

        security = DataFrame(list(self.securities_collection.find()))
        # security['Datetime'] = security['Datetime'].dt.tz_convert(None)
        # security.set_index('Symbol', inplace=True)
        # bars = bars[['Open', 'High', 'Low', 'Close', 'Volume', 'OpenInterest']]
        return security